package com.tamude.gpt.service.impl;

import com.alibaba.fastjson.JSONObject;
import com.tamude.gpt.service.ChatGptService;
import com.unfbx.chatgpt.utils.TikTokensUtil;
import constant.MessageRoleEnum;
import entity.ChatMessage;
import listener.OpenAISubscriber;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;
import util.PdfUtils;
import util.api.OpenAiWebClient;

import java.util.*;

import static constant.CommonConstant.MSG_CACHE;
import static constant.CommonConstant.PDF_CONTENT_CACHE;

@Slf4j
@Service
@RequiredArgsConstructor
public class ChatGptServiceImpl implements ChatGptService {

    private final OpenAiWebClient openAiWebClient;

    @Override
    public Flux<String> chatWithGpt(String msg, String pdfUrl, String userId) {

        String pdfId = pdfUrl.replace("https://pdf-container.oss-cn-beijing.aliyuncs.com/", "");

        //先判断pdf缓存中是否存在该pdf
        if (!PDF_CONTENT_CACHE.containsKey(pdfId)){
            StringBuilder sb = new StringBuilder();
            String pdfContent = PdfUtils.pdfToContent(pdfUrl);
            sb.append("假设你是一个假如你是一位科研工作者，阅读并理解了以下内容：\n")
                    .append(pdfContent)
                    .append("\n" + "请根据以上内容回答我的问题：\n");
            PDF_CONTENT_CACHE.put(pdfId, sb.toString());
        }

        //再判断用户是否有过历史记录
        if (MSG_CACHE.containsKey(userId)){
            //判断用户是否有针对该pdf的历史记录
            if (!MSG_CACHE.get(userId).containsKey(pdfId)){
                ChatMessage chatMessage = new ChatMessage(MessageRoleEnum.USER.getName(), PDF_CONTENT_CACHE.get(pdfId));
                MSG_CACHE.get(userId).put(pdfId, Collections.singletonList(chatMessage));
            }
        }else {
            Map<String,List<ChatMessage>> m = new HashMap<>();
            ChatMessage chatMessage = new ChatMessage(MessageRoleEnum.USER.getName(), PDF_CONTENT_CACHE.get(pdfId));
            m.put(pdfId,Collections.singletonList(chatMessage));
            MSG_CACHE.put(userId,m);
        }

        List<ChatMessage> historyMsgList = MSG_CACHE.get(userId).get(pdfId);

        int requestTokens = TikTokensUtil.tokens("gpt-3.5-turbo", msg);
        StringBuilder sb = new StringBuilder();

        return Flux.create(emitter -> {
            OpenAISubscriber subscriber = new OpenAISubscriber(emitter, userId, this, historyMsgList);
            Flux<String> openAiResponse =
                    openAiWebClient.getChatResponse(userId, msg, null, null, null);
            //遍历数据流返回的内容，计算tokens
            openAiResponse
                    .doOnNext(e -> {
                        if ("[DONE]".equals(e)){
                            return;
                        }
                        try {
                            String resContent = JSONObject.parseObject(e)
                                    .getJSONArray("choices")
                                    .getJSONObject(0)
                                    .getJSONObject("delta")
                                    .getString("content");
                            sb.append(resContent);
                        } catch (Exception ex) {
                            log.error("响应数据流解析失败！\n", ex);
                            throw new RuntimeException("响应数据流解析失败！" + ex);
                        }
                    })
                    .doOnComplete(() -> {
                        int responseTokens = TikTokensUtil.tokens("gpt-3.5-turbo", sb.toString());
                        log.info("请求消耗toekns:{}", requestTokens);
                        log.info("响应消耗tokens:{}", responseTokens);
                        log.info("总计消耗tokens:{}", responseTokens+requestTokens);
                        //历史记录写入
                        historyMsgList.add(new ChatMessage(MessageRoleEnum.ASSISTANT.getName(), sb.toString()));
                        MSG_CACHE.get(userId).put(pdfId,historyMsgList);
                    })
                    .subscribe(subscriber);
            emitter.onDispose(subscriber);
        });

    }

    @Override
    public void completed(List<ChatMessage> questions, String sessionId, String response) {

    }

    @Override
    public void fail(List<ChatMessage> questions, String sessionId, String errorMsg) {
        log.error("openai 处理失败 userId:{},errorMsg:{}", sessionId, errorMsg);
    }

    @Override
    public void clearHistory(String sessionId) {
    }
}
